import time
import numpy
from tqdm import tqdm
cimport numpy
cimport cython

DTYPE = numpy.float
ctypedef numpy.float_t DTYPE_t

@cython.boundscheck(False) # turn off bounds-checking for entire function
@cython.wraparound(False)  # turn off negative index wrapping for entire function
def generate_geo_index(numpy.ndarray[float, ndim=2] label_coords,
                       numpy.ndarray[float, ndim=3] coord_data, 
                       numpy.ndarray[int, ndim=2] patch_inds_x, 
                       numpy.ndarray[int, ndim=2] patch_inds_y, 
                       numpy.ndarray[float, ndim=1] label_depths, 
                       float lat_margin,
                       float lon_margin):

    cdef int N = label_coords.shape[0]
    y_inds, x_inds, depths = [], [], []
    cdef float lat_max = coord_data[:, :, 0].max()
    cdef float lat_min = coord_data[:, :, 0].min()
    cdef float lon_max = coord_data[:, :, 1].max()
    cdef float lon_min = coord_data[:, :, 1].min()
    cdef float lon, lat

    for i in tqdm(range(N)):
        lon, lat = label_coords[i][0], label_coords[i][1]
        if lat > lat_max or lat < lat_min or \
            lon > lon_max or lon < lon_min:
                continue

        y_ind, x_ind = numpy.where((coord_data[:, :, 0] > lat - lat_margin) & \
                                   (coord_data[:, :, 0] < lat + lat_margin) & \
                                   (coord_data[:, :, 1] > lon - lon_margin) & \
                                   (coord_data[:, :, 1] < lon + lon_margin))
        if y_ind.shape[0] > 0:
            y_inds.append(patch_inds_y + int(y_ind.mean()))
            x_inds.append(patch_inds_x + int(x_ind.mean()))
            depths.append(label_depths[i])

    return y_inds, x_inds, depths

